import datetime
import time

import torch
from torch.utils.data import DataLoader
from torchsummary import summary

from yolo import Model
from yolo_resnet34 import YOLOv1_resnet
from dataset import Banana, VOC2012
from loss import Loss_yolov1
from device import try_gpu

if __name__ == '__main__':
    epoch = 10
    batchsize = 5
    lr = 0.001

    train_data = Banana(is_train=True)
    # train_data = VOC2012(is_train=True, is_aug=True)
    train_dataloader = DataLoader(train_data, batch_size=batchsize, shuffle=True)

    net = Model()
    # net = YOLOv1_resnet()

    # for param in net.resnet.parameters():
    #     param.requires_grad = False

    # trainer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)
    optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.1, weight_decay=0.0005)
    # optimizer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=0.0005)

    criterion = Loss_yolov1()
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = try_gpu()
    # device = torch.device("cpu")

    net = net.to(device)

    # print(summary(net, (3, 448, 448)))

    start_time = time.time()
    for e in range(epoch):
        net.train()
        for i, (inputs, labels) in enumerate(train_dataloader):
            # inputs = inputs.cuda()
            inputs = inputs.to(device)
            # labels = labels.float().cuda()
            labels = labels.float().to(device)
            pred = net(inputs)
            loss = criterion(pred, labels)

            optimizer.zero_grad()  # 清除梯度
            loss.backward()  # 更新梯度
            optimizer.step()  # 更新权重

            # print("Epoch %d/%d| Step %d/%d| Loss: %.2f" % (e, epoch, i, len(train_data) // batchsize, loss))
        print("Epoch %d/%d | Loss: %.2f" % (e, epoch, loss))

        #     yl = yl + loss
        #     if is_vis and (i + 1) % 100 == 0:
        #         vis.line(np.array([yl.cpu().item() / (i + 1)]), np.array([i + e * len(train_data) // batchsize]),
        #                  win=viswin1, update='append')
        # if (e + 1) % 10 == 0:
        #     torch.save(model, "./models_pkl/YOLOv1_epoch" + str(e + 1) + ".pkl")
    end_time = time.time()
    speed = (end_time - start_time) / 60 / epoch
    print("speed:", speed, "min/epoch")
    torch.save(
        net.state_dict(),
        'YOLOV1'+datetime.datetime.now().strftime('%Y-%m-%d-%H_%M_%S')+'.pth'
    )
    # X = torch.normal(0, 1, (1, 3, 448, 448))
    # Y = net(X)
    #
    # print(Y)
    # print(Y.size())
    # print(Y.max())

